/*
 * Copyright 2010 Google Inc.
 * 
 * Licensed under the Apache License, Version 2.0 (the "License"); you may not
 * use this file except in compliance with the License. You may obtain a copy of
 * the License at
 * 
 * http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
 * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
 * License for the specific language governing permissions and limitations under
 * the License.
 */
package com.google.websocket;

import static com.google.websocket.Util.writeLength;

import org.eclipse.jetty.util.Utf8StringBuilder;

import java.io.BufferedInputStream;
import java.io.BufferedOutputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.io.OutputStreamWriter;
import java.io.PrintWriter;
import java.nio.charset.Charset;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

/**
 * A class representing a WebSocket connection.
 */
public class DraftWebSocket implements FragmentableWebSocket {

  /**
   * Simple callback that aggregates fragments.
   */
  public abstract class SimpleCallback implements FragmentableCallback {
    
    private int origOpcode;
    private List<byte[]> fragments = new ArrayList<byte[]>();
    private Set<ExtensionData> messageMetadata;

    public void onBinaryMessage(byte[] message,
        Set<ExtensionData> metadata) throws IOException {
      onUnexpectedMessage(FRAME_BINARY, fragments, metadata);
    }

    public void onClose() {
      // Do nothing by default
    }

    public void onFrameReceived(boolean more, int opcode,
        Set<ExtensionData> metadata, byte[] payload) throws IOException {
      if (opcode == DraftWebSocket.FRAME_CONTINUATION) {
        assert !fragments.isEmpty();
        assert metadata == null;
      } else {
        // TODO: per-frame metadata?  Aggregate fragmented metadata?
        messageMetadata = metadata;
        this.origOpcode = opcode;
      }
      fragments.add(payload);
      if (more) {
        return;
      }
      switch (origOpcode) {
        case DraftWebSocket.FRAME_BINARY:
          byte[] message = collectBytes();
          onBinaryMessage(message, messageMetadata);
          break;
        case DraftWebSocket.FRAME_TEXT:
          Utf8StringBuilder buf = new Utf8StringBuilder();
          for (byte[] fragment : fragments) {
            buf.append(fragment, 0, fragment.length);
          }
          onTextMessage(buf.toString(), messageMetadata);
          break;
        case DraftWebSocket.FRAME_CLOSE:
          onClose();
          DraftWebSocket.this.bufIn.close();
          break;
        case DraftWebSocket.FRAME_PING:
          byte[] ping = collectBytes();
          DraftWebSocket.this.sendFrame(false, DraftWebSocket.FRAME_PONG, (byte[]) null, ping);
          break;
        default:
          onUnexpectedMessage(origOpcode, fragments, messageMetadata);
          break;
      }
      fragments.clear();
    }

    public abstract void onFramingError();

    public void onTextMessage(String message,
        Set<ExtensionData> metadata) throws IOException {
      onUnexpectedMessage(FRAME_TEXT, fragments, metadata);
    }

    /**
     * Called when an unrecognized opcode is received.
     * 
     * @param opcode
     * @param fragments
     * @param metadata
     */
    public abstract void onUnexpectedMessage(int opcode, List<byte[]> fragments,
        Set<ExtensionData> metadata);

    private byte[] collectBytes() {
      int totalLength = 0;
      for (byte[] fragment : fragments) {
        totalLength += fragment.length;
      }
      byte[] message = new byte[totalLength];
      int ofs = 0;
      for (byte[] fragment : fragments) {
        System.arraycopy(fragment, 0, message, ofs, fragment.length);
        ofs += fragment.length;
      }
      return message;
    }
  }

  private class DelegatingCallback extends SimpleCallback {
    
    private final Callback callback;

    public DelegatingCallback(Callback callback) {
      this.callback = callback;
    }

    @Override
    public void onBinaryMessage(
        byte[] message, Set<ExtensionData> metadata)
        throws IOException {
      callback.onBinaryMessage(message, metadata);
    }

    @Override
    public void onClose() {
      callback.onClose();
    }

    @Override
    public void onFramingError() {
      callback.onFramingError();
    }

    @Override
    public void onTextMessage(String message,
        Set<ExtensionData> metadata) throws IOException {
      callback.onTextMessage(message, metadata);
    }

    @Override
    public void onUnexpectedMessage(int opcode, List<byte[]> fragments,
        Set<ExtensionData> metadata) {
      callback.onFramingError();
    }
  }

  // Integer constants are used rather than enums so clients of this library
  // can make use of new (or private-use) frame types before they are defined
  // here.
  public static final int FRAME_CONTINUATION = 0;

  public static final int FRAME_CLOSE = 1;

  public static final int FRAME_PING = 2;

  public static final int FRAME_PONG = 3;

  public static final int FRAME_TEXT = 4;

  public static final int FRAME_BINARY = 5;

  // flags and fields of the first header byte
  static final int MORE = 0x80;

  static final int OPCODE_MASK = 0x0f;

  // flags and fields of the second header byte
  static final int HAS_METADATA = 0x80;
  static final int LENGTH_MASK = 0x7f;
  static final int LENGTH_16 = 126;
  static final int LENGTH_63 = 127;

  private static final Charset UTF8 = Charset.forName("UTF-8");

  private static byte[] buildMetadataBytes(Set<ExtensionData> metadata) {
    if (metadata == null || metadata.isEmpty()) {
      return null;
    }
    int n = metadata.size();
    ByteArrayOutputStream bstr = new ByteArrayOutputStream();
    int i = 0;
    for (ExtensionData extensionData : metadata) {
      Extension extension = extensionData.getExtension();
      assert extension != null;
      assert extension.usesMetadata();
      int id = extension.getId();
      assert id >= 0 && id <= 255;
      byte[] bytes = extension.generateExtensionData(extensionData);
      boolean lastOne = (++i == n);
      try {
        Util.writeLength(bstr, bytes.length, !lastOne);
        bstr.write(id);
        bstr.write(bytes);
      } catch (IOException e) {
        throw new RuntimeException("Error writing to ByteArrayOutputStream", e);
      }
    }
    return bstr.toByteArray();
  }

  private final BufferedInputStream bufIn;

  private final BufferedOutputStream bufOut;

  private final Set<Extension> extensions;

  private final boolean supportsMetadata;

  /**
   * Create a WebSocket connection on an already-established connection.  The
   * caller is responsible for completing the connection establishment phase and
   * any extension negotiation before creating this object.
   * 
   * TODO: perhaps pass in a negotiated maximum frame size and fragment as
   * necessary to reflect that?  Or should that just be in extensions?
   * 
   * @param in
   * @param out
   * @param extensions extensions which were negotiated during the handshake,
   *     may be null
   */
  public DraftWebSocket(InputStream in, OutputStream out,
      Set<Extension> extensions) {
    bufIn = in instanceof BufferedInputStream ? (BufferedInputStream) in
        : new BufferedInputStream(in);
    bufOut = out instanceof BufferedOutputStream ? (BufferedOutputStream) out
        : new BufferedOutputStream(out);
    this.extensions = extensions;
    boolean tempSupportsMetadata = false;
    if (extensions != null) {
      for (Extension extension : extensions) {
        if (extension.usesMetadata()) {
          tempSupportsMetadata = true;
        }
      }
    }
    supportsMetadata = tempSupportsMetadata;
  }

  /**
   * Close the WebSocket.  Sends a graceful close frame to the remote endpoint,
   * and no further frames may be sent.
   * 
   * @throws IOException
   */
  public void close() throws IOException {
    sendFrame(false, FRAME_CLOSE, (byte[]) null, new byte[0]);
    bufOut.close();
  }

  /**
   * Read one frame and process it via the supplied callback.
   * 
   * @param callback
   * @return true if there are more frames in this message
   * @throws IOException
   */
  public boolean readOneFrame(FragmentableCallback callback)
      throws IOException {
    int hdrByte = Util.readByte(bufIn);
    boolean more = (hdrByte & MORE) != 0;
    int opcode = hdrByte & OPCODE_MASK;
    hdrByte = Util.readByte(bufIn);
    boolean hasMetadata = (hdrByte & HAS_METADATA) != 0;
    if (hasMetadata && !supportsMetadata) {
      callback.onFramingError();
      bufIn.close();
      throw new IOException("WebSocket framing error");
    }
    long longLength = Util.readLength(hdrByte, bufIn);
    Set<ExtensionData> metadata = null;
    if (hasMetadata) {
      metadata = new HashSet<ExtensionData>();
      longLength -= readMetadata(metadata); 
    }
    if (longLength > Integer.MAX_VALUE) {
      throw new IOException("Frame size too large");
    }
    int length = (int) longLength;
    byte[] payload = new byte[length];
    Util.readFully(payload, bufIn);
    callback.onFrameReceived(more, opcode, metadata, payload);
    return !more;
  }

  /**
   * Process a single message via the supplied callback.
   * 
   * @param callback
   * @throws IOException
   */
  public void readOneMessage(Callback callback) throws IOException {
    DelegatingCallback cb = new DelegatingCallback(callback);
    while (!readOneFrame(cb)) {
      // empty
    }
  }

  /**
   * Send an arbitrary stream of binary data as a binary messsage.
   *
   * @param metadata
   * @return an OutputStream to write the message payload to -- it must be
   *     closed before any other non-control messages are written to this
   *     WebSocket 
   * @throws IOException
   */
  public OutputStream sendBinaryMessage(Set<ExtensionData> metadata)
      throws IOException {
    return sendStream(FRAME_BINARY, metadata);
  }

  /**
   * Send a binary message.
   *
   * @param metadata
   * @param message
   * @throws IOException
   */
  public void sendBinaryMessage(Set<ExtensionData> metadata, byte[] message)
      throws IOException {
    sendFrame(false, FRAME_BINARY, metadata, message);
  }

  /**
   * Send an arbitrary WebSocket frame.
   * 
   * @param more true if this is not the last frame in this message
   * @param opcode
   * @param metadata
   * @param payload
   * @throws IOException
   */
  public void sendFrame(boolean more, int opcode,
      Set<ExtensionData> metadata, byte[] payload) throws IOException {
    sendFrame(more, opcode, buildMetadataBytes(metadata), payload);
  }

  /**
   * Send an arbitrary message type from a stream.
   *
   * @param opcode
   * @param metadata
   * @return an OutputStream to write the message payload to -- it must be
   *     closed before any other non-control messages are written to this
   *     WebSocket 
   * @throws IOException
   */
  public OutputStream sendStream(final int opcode,
      final Set<ExtensionData> metadata) throws IOException {
    return new OutputStream() {

      private static final int BUF_SIZE = 8192;

      // the maximum frame header size - this is the offset into the buffer
      // where user data is written to allow the header to be inserted before it
      private static final int MAX_HEADER_SIZE = 10; 

      private final byte[] buf = new byte[BUF_SIZE]; 
      private int bufIdx;

      private int currentOpcode = opcode;
      private boolean hasMetadata = false;

      {
        // write metadata, and position bufIdx for payload data
        bufIdx = MAX_HEADER_SIZE;
        byte[] metadataBytes = buildMetadataBytes(metadata);
        if (metadataBytes != null) {
          int metadataLength = metadataBytes.length;
          // special-case the metadata being larger than our buffer size
          // TODO: fragment the metadata
          if (metadataLength >= BUF_SIZE) {
            int hdrLen = buildFrameHeader(true, true, metadataLength);
            bufOut.write(buf, MAX_HEADER_SIZE - hdrLen, hdrLen);
            bufOut.write(metadataBytes);
            bufOut.flush();
            currentOpcode = FRAME_CONTINUATION;
          } else {
            System.arraycopy(metadataBytes, 0, buf, bufIdx, metadataLength);
            bufIdx += metadataLength;
            hasMetadata = true;
          }
        }
      }

      @Override
      public void close() throws IOException {
        // on stream close, write the remainder of the buffer as the last frame
        writeFrameFromBuffer(false);
      }

      @Override
      public void flush() throws IOException {
        writeFrameFromBuffer(true);
      }

      @Override
      public void write(int b) throws IOException {
        if (bufIdx >= BUF_SIZE) {
          writeFrameFromBuffer(true);
        }
        buf[bufIdx++] = (byte) b;
      }

      /**
       * Generate the frame header before the start of user data.
       *
       * @param more true if the MORE flag should be set
       * @param hasMetadata true if the HAS_METADATA flag should be set
       * @param len total length of metadata and payload in this frame
       * @return length of generated header
       */
      private int buildFrameHeader(boolean more, boolean hasMetadata, int len) {
        int hdrLen;
        int lengthByte = hasMetadata ? HAS_METADATA : 0;
        if (len < LENGTH_16) {
          hdrLen = 2;
          buf[MAX_HEADER_SIZE - 1] = (byte) (len | lengthByte);
        } else if (len < 65536) {
          hdrLen = 4;
          buf[MAX_HEADER_SIZE - 1] = (byte) (len & 255);
          buf[MAX_HEADER_SIZE - 2] = (byte) ((len >> 8) & 255);
          buf[MAX_HEADER_SIZE - 3] = (byte) (lengthByte | LENGTH_16);
        } else {
          hdrLen = 10;
          buf[MAX_HEADER_SIZE - 1] = (byte) (len & 255);
          buf[MAX_HEADER_SIZE - 2] = (byte) ((len >> 8) & 255);
          buf[MAX_HEADER_SIZE - 3] = (byte) ((len >> 16) & 255);
          buf[MAX_HEADER_SIZE - 4] = (byte) ((len >> 24) & 255);
          buf[MAX_HEADER_SIZE - 5] = (byte) 0;
          buf[MAX_HEADER_SIZE - 6] = (byte) 0;
          buf[MAX_HEADER_SIZE - 7] = (byte) 0;
          buf[MAX_HEADER_SIZE - 8] = (byte) 0;
          buf[MAX_HEADER_SIZE - 9] = (byte) (lengthByte | LENGTH_63);
        }
        buf[MAX_HEADER_SIZE - hdrLen] = (byte) ((more ? MORE : 0)
            | currentOpcode);
        return hdrLen;
      }

      /**
       * Write a full frame.
       *
       * @param more
       * @throws IOException
       */
      private void writeFrameFromBuffer(boolean more) throws IOException {
        int hdrLen = buildFrameHeader(more, hasMetadata,
            bufIdx - MAX_HEADER_SIZE);
        bufOut.write(buf, MAX_HEADER_SIZE - hdrLen, bufIdx - MAX_HEADER_SIZE
            + hdrLen);
        bufOut.flush();
        currentOpcode = FRAME_CONTINUATION;
        bufIdx = MAX_HEADER_SIZE;
        hasMetadata = false;
      }
    };
  }

  /**
   * Write an arbitrarily long text message.
   * 
   * @param metadata
   * @return PrintWriter to write message on; must be closed before sending any
   *     other non-control messages on this WebSocket
   * @throws IOException
   */
  public PrintWriter sendTextMessage(Set<ExtensionData> metadata)
      throws IOException {
    OutputStream stream = sendStream(FRAME_TEXT, metadata);
    return new PrintWriter(new OutputStreamWriter(stream, UTF8));
  }

  /**
   * Write a text message.
   *
   * @param metadata
   * @param message
   * @throws IOException
   */
  public void sendTextMessage(Set<ExtensionData> metadata, String message)
      throws IOException {
    byte[] bytes = message.getBytes(UTF8);
    sendFrame(false, FRAME_TEXT, metadata, bytes);
  }

  /**
   * Extracts metadata from a WebSocket frame.
   * 
   * TODO: handle fragmented metadata
   * 
   * @param metadata
   * @return number of bytes read
   * @throws IOException
   */
  private long readMetadata(Set<ExtensionData> metadata) throws IOException {
    int lengthByte = Util.readByte(bufIn);
    long bytesRead = 0;
    do {
      long metaLengthLong = Util.readLength(lengthByte, bufIn);
      bytesRead += Util.numberOfLengthBytes(metaLengthLong);
      int idByte = Util.readByte(bufIn);
      bytesRead++;
      if (metaLengthLong > Integer.MAX_VALUE) {
        throw new IOException("metadata too large");
      }
      int metaLength = (int) metaLengthLong;
      byte[] metadataValue = new byte[metaLength];
      Util.readFully(metadataValue, bufIn);
      bytesRead += metaLength;
      Extension extension = ExtensionRegistry.findExtensionByMetadataId(idByte);
      assert extensions.contains(extension);
      ExtensionData extensionData = extension.parseExtensionData(metadataValue);
      metadata.add(extensionData);
    } while ((lengthByte & MORE) != 0);
    return bytesRead;
  }

  /**
   * Send a frame.
   * 
   * @param more true if more frames of this message follow
   * @param opcode frame opcode to use
   * @param payload payload bytes
   * @throws IOException
   */
  private void sendFrame(boolean more, int opcode, byte[] metadataBytes,
      byte[] payload) throws IOException {
    bufOut.write((more ? MORE : 0) | opcode);
    long length = payload.length;
    if (metadataBytes != null) {
      length += metadataBytes.length;
    }
    writeLength(bufOut, length, metadataBytes != null);
    if (metadataBytes != null) {
      bufOut.write(metadataBytes);
    }
    bufOut.write(payload);
    bufOut.flush();
  }
}
